from PyGMO.problem import base_stochastic
from numpy import sqrt, linspace, dot, complex128, ones, reshape,\
                  array, eye, kron, append, zeros
from time import time
from utils import pulse_fitness #NEVER import prep_dict_dep

def constant_pulse(params, tvec):
    if len(params) != 1:
        raise ValueError("constant pulses only take iterables of one argument")
    return params[0] * ones(len(tvec))

class PulseProblem(base_stochastic):
    """
    This class encodes pulse optimization problems for solution by
    PyGMO. Each instance of the class takes as input:

     + A function `pulse_func` which maps the optimization parameters
       and a time vector onto an amplitude vector for the pulse.

     + Absolute bounds on the pulse parameters.

     + An optional T_1 decay rate.

     + An optional T_2 decay rate.

     + A fixed drive time, which is not used except when the decay
       rates are 0.

     + A quantum efficiency eta
    """
    def __init__(self, pulse_func=constant_pulse, bounds=(0.4,0.7), gamma_d=0., gamma_r=0.,
                    fixed_time=10., eta=0.5):
        #Sanitize input:
        """
        n_pulse_args = pulse_func.func_code.co_argcount
        has_decoherence = int((gamma_d != 0.) or (gamma_r != 0.))
        if len(bounds) != n_pulse_args + has_decoherence:
                raise ValueError("There must be as many tuples of " +
                    "bounds as there are arguments to the pulse " + 
                        "function, plus one if there's decoherence.")
        """
        super(PulseProblem, self).__init__(len(bounds), 0)

        self.set_bounds(0.4,0.7)

        #Private Members
        self._pulse_func = pulse_func
        self._gamma_d = gamma_d
        self._gamma_r = gamma_r
        self._fixed_time = fixed_time
        self._eta = eta
        self._input_dict = prep_dict_indep()

    def _objfun_impl(self, params):
        from c_traj import get_trajectory

        input_dict = dict(self._input_dict.items() + 
                                prep_dict_dep(self, params).items())
        
        return pulse_fitness(input_dict)
    
    # reimplement the virtual method that compares fitnesses
    def _compare_fitness_impl(self, f1, f2):
        return f1[0] > f2[0] #maximization, not minimization

class TestProblem(base_stochastic):
    """docstring for TestProblem"""
    def __init__(self, bounds=(0.0,20.0)):
        super(TestProblem, self).__init__(len(bounds),0)
        
        self.set_bounds(0.0,20.0)

    def _compare_fitness_impl(self, f1, f2):
        return f1[0] > f2[0] #maximization, not minimization    

    def _objfun_impl(self, params):
        from numpy.random import randn
        return (5. - (params[0] - 10.)**2 + 0.1 * randn(),)


def prep_dict_indep():
    """
    Place to dump messy code which sets the problem-independent values 
    of the input_dict. This ensures that unnecessary kronecker products
    are only evaluated at initialization, not every time the objective
    function is called.
    """

    initstate = 1./sqrt(8.) * ones((8, 1), dtype='complex128')
    initMat = initstate * initstate.T
    init = reshape(initMat.T, 64, 1)

    #Time-independent part of Liouvillian
    sigmaZ = array([[1, 0], [0, - 1]])
    sigmaM = array([[0, 1], [0, 0]])
    smalleye = eye(2)
    fulleye = eye(8)

    # The dephasing operator for qubit 1
    Ad1 = reduce(kron,(sigmaZ, smalleye, smalleye))
    Ld1 = kron(Ad1, Ad1.T) - dot(0.5, kron(dot(
        Ad1.T, Ad1), fulleye)) - dot(0.5, kron(fulleye, dot(Ad1.T, Ad1).T))

    # The dephasing operator for qubit 2
    Ad2 = reduce(kron,(smalleye, sigmaZ, smalleye))
    Ld2 = kron(Ad2, Ad2.T) - dot(0.5, kron(dot(
        Ad2.T, Ad2), fulleye)) - dot(0.5, kron(fulleye, dot(Ad2.T, Ad2).T))

    # The dephasing operator for qubit 3
    Ad3 = reduce(kron,(smalleye, smalleye, sigmaZ))
    Ld3 = kron(Ad3, Ad3.T) - dot(0.5, kron(dot(
        Ad3.T, Ad3), fulleye)) - dot(0.5, kron(fulleye, dot(Ad3.T, Ad3).T))

    # The relaxation operator for qubit 1
    Ar1 = reduce(kron,(sigmaM, smalleye, smalleye))
    Ar1dagger = Ar1.T
    Lr1 = kron(Ar1, Ar1dagger.T) - dot(0.5, kron(dot(
        Ar1dagger, Ar1), fulleye)) - dot(0.5, kron(fulleye, dot(Ar1dagger, Ar1).T))

    # The relaxation operator for qubit 2
    Ar2 = reduce(kron,(smalleye, sigmaM, smalleye))
    Ar2dagger = Ar2.T
    Lr2 = kron(Ar2, Ar2dagger.T) - dot(0.5, kron(dot(
        Ar2dagger, Ar2), fulleye)) - dot(0.5, kron(fulleye, dot(Ar2dagger, Ar2).T))

    # The relaxation operator for qubit 3
    Ar3 = reduce(kron,(smalleye, smalleye, sigmaM))
    Ar3dagger = Ar3.T
    Lr3 = kron(Ar3, Ar3dagger.T) - dot(0.5, kron(dot(
        Ar3dagger, Ar3), fulleye)) - dot(0.5, kron(fulleye, dot(Ar3dagger, Ar3).T))

    #Ideal Plus/Minus density matrices
    ideal_p = 0.5 * array([1,0,0,1,0,1,1,0], dtype=complex128).T
    ideal_p = kron(ideal_p, ideal_p)
    ideal_m = 0.5 * array([0,1,1,0,1,0,0,1], dtype=complex128).T
    ideal_m = kron(ideal_m, ideal_m)

    indep_dict = {'dim': 8, 'chi': 1., 'kappa1': 2., 'kappa2': 2., 
                    'Delta_a': sqrt(3.), 'Delta_b': -sqrt(3.),
                    'dt': 0.005, 'init': init, 'ideal_p': ideal_p,
                    'ideal_m': ideal_m, 'Ld1': Ld1, 'Ld2': Ld2,
                    'Ld3': Ld3, 'Lr1': Lr1, 'Lr2': Lr2, 'Lr3': Lr3}

    return indep_dict

def prep_dict_dep(probject, params):
    """
    Prepares a dictionary for input to get_trajectory, based on the
    parameters of the PulseProblem object.

    This function takes the time to be optimized to be the drive
    time, so that the measurement turns off at the same time as the
    drive, and the simulation continues for another 2/chi to capture
    the ring-down. 
    """
    
    if (probject._gamma_d == 0) and (probject._gamma_r == 0):
        t_drive = probject._fixed_time
        pass_params = params
    else:
        t_drive = params[-1]
        pass_params = params[:-1]
    
    tmeas_off = t_drive

    #Decay rates, given by master decay rates gamma_d and gamma_r

    #input_dict['T'] = distortion_fudge_factor * input_dict['tmeas_off'] # final time
    T = tmeas_off # only drive time here

    #Cavity input, given by pulse_func:
    dt = probject._input_dict['dt']
    N = int(round(T / dt))
    tvec = linspace(0., T, N)
    epsilon = probject._pulse_func(pass_params, tvec)

    # The time independent deterministic part of Liouvillian, affected by kappa
    #input_dict['D1'] = (dot(dot(0.5, input_dict['gamma_d1']), Ld1) + \
    #                    dot(dot(0.5, input_dict['gamma_d2']), Ld2) + \
    #                    dot(dot(0.5, input_dict['gamma_d3']), Ld3) + \
    #                    dot(dot(0, (dot((g / Delta_q) ** 2, (input_dict['kappa1'] + input_dict['kappa2'])) + \
    #                        input_dict['gamma_1'])), (Lr1 + Lr2 + Lr3))).astype(complex128)
    
    gamma_d = probject._gamma_d
    gamma_r = probject._gamma_r
    Ld1 = probject._input_dict['Ld1']; Ld2 = probject._input_dict['Ld2']
    Ld3 = probject._input_dict['Ld3']
    Lr1 = probject._input_dict['Lr1']; Lr2 = probject._input_dict['Lr2']
    Lr3 = probject._input_dict['Lr3']
    kappa1 = probject._input_dict['kappa1']; kappa2 = probject._input_dict['kappa2']
    D1 = (0.5 * gamma_d * Ld1 + 0.5 * gamma_d * Ld2 + 0.5 * gamma_d * Ld3 + 
            0 * (1./1600. * (kappa1 + kappa2) + gamma_r ) * (Lr1 + Lr2 + Lr3)).astype(complex128)

    dep_dict = {'D1': D1, 'eta': probject._eta, 'tmeas_on': 0.,
                'tmeas_off': tmeas_off, 'epsilon': epsilon, 'N': N, 'T': T}

    return dep_dict
